package cn.itcast.config;

import cn.itcast.constants.Constant;
import cn.itcast.tools.PackageTools;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.PromptChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SafeGuardAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class SpringAIConfig {

    @Bean
    public ChatClient chatClient(ChatClient.Builder builder,
                                 Advisor simpleLoggerAdvisor,
                                 Advisor messageChatMemoryAdvisor,
                                 Advisor promptChatMemoryAdvisor,
                                 Advisor safeGuardAdvisor,
                                 PackageTools packageTools
    ) {
        return builder
                .defaultSystem(Constant.SYSTEM_ROLE)
                .defaultAdvisors(simpleLoggerAdvisor,
                    promptChatMemoryAdvisor,
                    safeGuardAdvisor)
                .defaultTools(packageTools)
                .build();
    }

    @Bean
    public ChatMemory chatMemory() {
        return new InMemoryChatMemory();
    }

    @Bean
    public Advisor simpleLoggerAdvisor() {
        return new SimpleLoggerAdvisor();
    }

    @Bean
    public Advisor messageChatMemoryAdvisor(ChatMemory chatMemory) {
        return new MessageChatMemoryAdvisor(chatMemory);
    }

    @Bean
    public Advisor promptChatMemoryAdvisor(ChatMemory chatMemory) {
        return new PromptChatMemoryAdvisor(chatMemory);
    }

    @Bean
    public Advisor safeGuardAdvisor() {
        List<String> sensitiveWords = List.of("ikun", "哥哥");
        return new SafeGuardAdvisor(sensitiveWords,"该输入包含敏感词，不允许继续。",Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
    }

    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel).build();
    }
}
